//      distnmf_main1.cc
//      main file for Distributed Rows NMF
//        -- generate random NMF problems
//        -- Block coordinate descent (batch - in one node)
//        -- Master/Slave concensus (row-distributed)
//        -- TODO: Decentralized NMF (with bipartie graph of worker nodes)
//      
//      Copyright 2012 tqlong <tqlong@espada>
//      
//      This program is free software; you can redistribute it and/or modify
//      it under the terms of the GNU General Public License as published by
//      the Free Software Foundation; either version 2 of the License, or
//      (at your option) any later version.
//      
//      This program is distributed in the hope that it will be useful,
//      but WITHOUT ANY WARRANTY; without even the implied warranty of
//      MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
//      GNU General Public License for more details.
//      
//      You should have received a copy of the GNU General Public License
//      along with this program; if not, write to the Free Software
//      Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
//      MA 02110-1301, USA.

/* example run
generate random A=W*H: 
   ./distnmf --rand --m=10 --k=2 --n=5 --output=d1 --verbose --binary
read matrix
   ./distnmf --read --input=d1 --binary
 */

#include "helper/stl_helper.h"
#include "helper/arma_helper.h"
#include "helper/mpi_helper.h"

#include "nmf.h"

po::options_description desc("Allowed options");
po::variables_map vm;
mpi::communicator* p_world = NULL;

void process_options(int argc, char** argv);
void usage(int exitcode = 0);

void read_data();
void rand_data();
void test_mpi();

boost::posix_time::ptime time_start_all(boost::posix_time::microsec_clock::local_time());
boost::posix_time::ptime time_start(boost::posix_time::microsec_clock::local_time());

int main(int argc, char** argv)
{
  mpi::environment env(argc, argv);
  mpi::communicator world;
  p_world = &world;
  MSG("rank=" << p_world->rank() << " size=" << p_world->size());

  process_options(argc, argv);

  if (vm.count("help"))
    usage();
  else if (vm.count("read"))
    read_data();
  else if (vm.count("rand"))
    rand_data();
  else if (vm.count("bcd"))
    nmf_bcd();
  else if (vm.count("mpi"))
    test_mpi();
  else if (vm.count("ms"))
    master_slave_nmf();
  else if (vm.count("dectr"))
    decentralized_nmf();
  else
    usage();

  MSG("Duration = " << duration_all());
}

void process_options(int argc, char** argv)
{
  desc.add_options()
    ("help", "produce this usage")
    ("input", po::value<std::string>(), "input file")
    ("output", po::value<std::string>(), "output file")
    ("log", po::value<std::string>(), "log file")
    ("verbose", "print debug information")

    // RAND
    ("m", po::value<int>(), "number of rows of A")
    ("n", po::value<int>(), "number of cols of A")    
    ("k", po::value<int>(), "true dimension of W and H")
    ("s", po::value<double>()->default_value(0.5), "sparsity level of W and H")
    ("binary", "save output in binary format")

    // NMF
    ("alpha", po::value<double>(), "regularization parameter \alpha*|W|_F^2 + \alpha*|H|_F^2")
    ("rho", po::value<double>(), "augmented pernalty term \alpha*|W|_F^2 + \alpha*|H|_F^2")
    ("iter", po::value<int>(), "maximum number of iterations")
    ("graph", po::value<std::string>(), "graph file")

    //JOB TYPE
    ("read", "read data file\n--input [--output]")
    ("rand", "random data file A=WH (size m x k x n)\n--m --n --k --output [--binary]")
    ("bcd", "Block Coordinate Descent\n--input [--binary] --k --alpha --output --log")
    ("ms", "Master/Slave Concensus NMF\n--input [--binary] --k --alpha --output --log")
    ("dectr", "Decentralized NMF\n--input [--binary] --k --alpha --graph --output --log")    
    ("mpi", "Test MPI")
    ;
  po::store(po::parse_command_line(argc, argv, desc), vm);
  po::notify(vm);
}

void usage(int exitcode)
{
  std::cout << desc << std::endl;
  exit(exitcode);
}

void read_data()
{
  std::string inFile = VM(vm, "input", std::string);
  matrix_t A;
  read_sparse_matrix<matrix_t>(inFile, A, vm.count("binary"));
}

void rand_data()
{
  std::string outFile = VM(vm, "output", std::string);
  int m = VM(vm, "m", int), 
    n = VM(vm, "n", int), 
    k = VM(vm, "k", int);
  double sparsity = VM(vm, "s", double);
  
  // generate random W and H
  matrix_t W, H;
  gen_sparse_matrix<matrix_t>(W, m, k, sparsity);
  gen_sparse_matrix<matrix_t>(H, k, n, sparsity);

  // A = W*H
  matrix_t A = W*H;// + norm_rnd*s_A;
  double s_A = arma::accu(A)/(A.n_rows * A.n_cols)/100;
  A += arma::randn<matrix_t>(m, n) * s_A;
  for (uint i = 0; i < A.n_rows; i++)
    for (uint j = 0; j < A.n_cols; j++)
      if (A(i,j) <= 0) A(i,j) = 0;

  MSG_VERBOSE(vm, "W=" << W << ENDL << "H=" << H << ENDL << "A=" << A);

  // save to output file
  save_sparse_matrix<matrix_t>(outFile, A, vm.count("binary"));
}

void test_mpi()
{
  vec_t z;
  if (p_world->rank() == 0) {
    z = arma::randu<vec_t>(5);
  }
  mpi::broadcast(*p_world, z, 0);
  MSG("rank= "<< p_world->rank() << " z=" << z);

  matrix_t M;
  if (p_world->rank() == 0) {
    M = arma::randu<matrix_t>(2,2);
  }
  mpi::broadcast(*p_world, M, 0);
  MSG("rank= "<< p_world->rank() << " M=" << M);
}

boost::posix_time::ptime now() 
{
  return boost::posix_time::microsec_clock::local_time();
}

boost::posix_time::time_duration duration() 
{
  boost::posix_time::time_duration d = now() - time_start;
  time_start = now();
  return d;
}

boost::posix_time::time_duration duration_all() 
{
  return now() - time_start_all;
}

